"""
OpenAI Compatible API 提供商

支持所有兼容 OpenAI API 格式的大模型服务
"""

import json
import time
import os
from typing import Dict, Any, List, Optional, Union
import requests
from dataclasses import dataclass

from ..core.data import Extraction, ExampleData, ExtractResult, generate_extraction_schema


@dataclass
class ProviderConfig:
    """提供商配置"""
    base_url: str
    model_name: str
    api_key: str
    max_tokens: int = 4096
    temperature: float = 0.1
    timeout: int = 30
    max_retries: int = 3
    supports_function_calling: bool = True


class OpenAICompatibleProvider:
    """OpenAI Compatible API 提供商
    
    支持所有兼容 OpenAI Chat Completions API 格式的服务
    """
    
    def __init__(self, config: ProviderConfig):
        """初始化提供商
        
        Args:
            config: 提供商配置
        """
        self.config = config
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {config.api_key}",
            "Content-Type": "application/json"
        })
    
    def _build_messages(
        self, 
        text: str, 
        prompt_description: str, 
        examples: List[ExampleData]
    ) -> List[Dict[str, str]]:
        """构建消息列表
        
        Args:
            text: 要提取的文本
            prompt_description: 提取任务描述
            examples: 示例数据
            
        Returns:
            消息列表
        """
        messages = []
        
        # 系统消息
        system_prompt = f"""你是一个专业的信息提取助手。请根据以下要求从文本中提取结构化信息：

{prompt_description}

请严格按照以下格式返回结果，确保提取的文本与原文完全一致：
- extraction_class: 提取类别
- extraction_text: 从原文中提取的确切文本
- attributes: 相关属性（可选）

重要提示：
1. 只提取原文中明确存在的信息
2. 提取的文本必须与原文完全一致，不要改写或总结
3. 如果没有找到相关信息，返回空列表
4. 属性值应该简洁明确"""

        messages.append({"role": "system", "content": system_prompt})
        
        # 添加示例
        for i, example in enumerate(examples):
            messages.append({
                "role": "user", 
                "content": f"示例{i+1}文本：\n{example.text}"
            })
            
            example_extractions = []
            for ext in example.extractions:
                ext_dict = {
                    "extraction_class": ext.extraction_class,
                    "extraction_text": ext.extraction_text
                }
                if ext.attributes:
                    ext_dict["attributes"] = ext.attributes
                example_extractions.append(ext_dict)
            
            messages.append({
                "role": "assistant",
                "content": json.dumps({
                    "extractions": example_extractions
                }, ensure_ascii=False, indent=2)
            })
        
        # 实际要处理的文本
        messages.append({
            "role": "user",
            "content": f"请从以下文本中提取信息：\n{text}"
        })
        
        return messages
    
    def _build_function_calling_request(
        self,
        text: str,
        prompt_description: str, 
        examples: List[ExampleData]
    ) -> Dict[str, Any]:
        """构建 Function Calling 请求
        
        Args:
            text: 要提取的文本
            prompt_description: 提取任务描述
            examples: 示例数据
            
        Returns:
            请求数据
        """
        messages = self._build_messages(text, prompt_description, examples)
        
        # 生成 function schema
        schema = generate_extraction_schema(examples)
        
        function_def = {
            "name": "extract_information",
            "description": "从文本中提取结构化信息",
            "parameters": schema
        }
        
        return {
            "model": self.config.model_name,
            "messages": messages,
            "functions": [function_def],
            "function_call": {"name": "extract_information"},
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens
        }
    
    def _build_regular_request(
        self,
        text: str,
        prompt_description: str,
        examples: List[ExampleData]
    ) -> Dict[str, Any]:
        """构建常规请求（不使用 Function Calling）
        
        Args:
            text: 要提取的文本
            prompt_description: 提取任务描述
            examples: 示例数据
            
        Returns:
            请求数据
        """
        messages = self._build_messages(text, prompt_description, examples)
        
        return {
            "model": self.config.model_name,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens
        }
    
    def _make_request(self, request_data: Dict[str, Any]) -> Dict[str, Any]:
        """发送 API 请求
        
        Args:
            request_data: 请求数据
            
        Returns:
            响应数据
            
        Raises:
            Exception: 请求失败时抛出异常
        """
        url = f"{self.config.base_url.rstrip('/')}/chat/completions"
        
        for attempt in range(self.config.max_retries):
            try:
                response = self.session.post(
                    url,
                    json=request_data,
                    timeout=self.config.timeout
                )
                
                if response.status_code == 200:
                    return response.json()
                else:
                    error_msg = f"API请求失败: {response.status_code}"
                    try:
                        error_detail = response.json()
                        error_msg += f" - {error_detail}"
                    except:
                        error_msg += f" - {response.text}"
                    
                    if attempt == self.config.max_retries - 1:
                        raise Exception(error_msg)
                    else:
                        time.sleep(2 ** attempt)  # 指数退避
                        
            except requests.exceptions.RequestException as e:
                if attempt == self.config.max_retries - 1:
                    raise Exception(f"网络请求失败: {str(e)}")
                else:
                    time.sleep(2 ** attempt)
        
        raise Exception("请求失败，已达到最大重试次数")
    
    def _parse_response(self, response_data: Dict[str, Any]) -> List[Extraction]:
        """解析 API 响应
        
        Args:
            response_data: API 响应数据
            
        Returns:
            提取结果列表
        """
        try:
            choice = response_data["choices"][0]
            
            # 检查是否使用了 function calling
            if "function_call" in choice["message"]:
                # Function calling 响应
                function_call = choice["message"]["function_call"]
                arguments = json.loads(function_call["arguments"])
                extractions_data = arguments.get("extractions", [])
            else:
                # 常规响应，尝试解析 JSON
                content = choice["message"]["content"].strip()
                
                # 尝试提取 JSON 部分
                if "```json" in content:
                    start = content.find("```json") + 7
                    end = content.find("```", start)
                    json_str = content[start:end].strip()
                elif content.startswith("{"):
                    json_str = content
                else:
                    # 尝试找到第一个 { 和最后一个 }
                    start = content.find("{")
                    end = content.rfind("}") + 1
                    if start != -1 and end > start:
                        json_str = content[start:end]
                    else:
                        raise ValueError("无法在响应中找到有效的JSON")
                
                parsed_data = json.loads(json_str)
                extractions_data = parsed_data.get("extractions", [])
            
            # 转换为 Extraction 对象
            extractions = []
            for ext_data in extractions_data:
                extraction = Extraction(
                    extraction_class=ext_data["extraction_class"],
                    extraction_text=ext_data["extraction_text"],
                    attributes=ext_data.get("attributes")
                )
                extractions.append(extraction)
            
            return extractions
            
        except (KeyError, json.JSONDecodeError, ValueError) as e:
            raise Exception(f"解析响应失败: {str(e)}")
    
    def extract(
        self,
        text: str,
        prompt_description: str,
        examples: List[ExampleData],
        use_function_calling: Optional[bool] = None
    ) -> ExtractResult:
        """执行信息提取
        
        Args:
            text: 要提取的文本
            prompt_description: 提取任务描述
            examples: 示例数据
            use_function_calling: 是否使用 Function Calling，None 表示自动判断
            
        Returns:
            提取结果
        """
        start_time = time.time()
        
        try:
            # 决定是否使用 Function Calling
            if use_function_calling is None:
                use_function_calling = self.config.supports_function_calling
            
            # 构建请求
            if use_function_calling:
                request_data = self._build_function_calling_request(
                    text, prompt_description, examples
                )
            else:
                request_data = self._build_regular_request(
                    text, prompt_description, examples
                )
            
            # 发送请求
            response_data = self._make_request(request_data)
            
            # 解析响应
            extractions = self._parse_response(response_data)
            
            # 构建结果
            end_time = time.time()
            metadata = {
                "model_name": self.config.model_name,
                "provider": "openai_compatible",
                "processing_time": end_time - start_time,
                "use_function_calling": use_function_calling,
                "total_extractions": len(extractions)
            }
            
            return ExtractResult(
                text=text,
                extractions=extractions,
                metadata=metadata,
                success=True
            )
            
        except Exception as e:
            end_time = time.time()
            metadata = {
                "model_name": self.config.model_name,
                "provider": "openai_compatible", 
                "processing_time": end_time - start_time,
                "error": str(e)
            }
            
            return ExtractResult(
                text=text,
                extractions=[],
                metadata=metadata,
                success=False,
                error_message=str(e)
            )


def create_provider_from_preset(
    provider_name: str,
    api_key: Optional[str] = None,
    **kwargs
) -> OpenAICompatibleProvider:
    """从预设配置创建提供商
    
    Args:
        provider_name: 预设提供商名称
        api_key: API 密钥，如果不提供则从环境变量获取
        **kwargs: 额外配置参数
        
    Returns:
        提供商实例
    """
    from .presets import get_preset_config, get_env_var_name
    
    preset = get_preset_config(provider_name)
    if not preset:
        raise ValueError(f"未知的预设提供商: {provider_name}")
    
    # 获取 API 密钥
    if not api_key:
        env_var = get_env_var_name(provider_name)
        if env_var:
            api_key = os.getenv(env_var)
        
        if not api_key:
            raise ValueError(f"未找到 API 密钥，请设置环境变量 {env_var} 或直接传入 api_key 参数")
    
    # 合并配置
    config_dict = {
        "base_url": preset["base_url"],
        "model_name": preset["model_name"],
        "api_key": api_key,
        "max_tokens": preset.get("max_tokens", 4096),
        "supports_function_calling": preset.get("supports_function_calling", True)
    }
    config_dict.update(kwargs)
    
    config = ProviderConfig(**config_dict)
    return OpenAICompatibleProvider(config)


def create_custom_provider(
    base_url: str,
    model_name: str,
    api_key: str,
    **kwargs
) -> OpenAICompatibleProvider:
    """创建自定义提供商
    
    Args:
        base_url: API 基础 URL
        model_name: 模型名称
        api_key: API 密钥
        **kwargs: 额外配置参数
        
    Returns:
        提供商实例
    """
    config_dict = {
        "base_url": base_url,
        "model_name": model_name,
        "api_key": api_key
    }
    config_dict.update(kwargs)
    
    config = ProviderConfig(**config_dict)
    return OpenAICompatibleProvider(config)
